from helper import *
from query import *
import numpy as np
import torch
from deepctr_torch.inputs import SparseFeat, DenseFeat, VarLenSparseFeat, get_feature_names
from deepctr_torch.models import DIEN
import os
os.environ["CUDA_VISIBLE_DEVICES"] = "7"

def get_fd():
    feature_columns = [
                    SparseFeat('pos_custom_id',1371980,embedding_dim=64,use_hash=False),
                    SparseFeat('pos_Active',2,embedding_dim=4,use_hash=False),
                    SparseFeat('pos_FN',2,embedding_dim=4,use_hash=False),
                    SparseFeat('pos_club_member_status',2,embedding_dim=4,use_hash=False),
                    SparseFeat('pos_fashion_news_frequency',2,embedding_dim=4,use_hash=False),
                    SparseFeat('pos_artid',105542+1,embedding_dim=64,use_hash=False),
                    SparseFeat('pos_pcode',105542+1,embedding_dim=64,use_hash=False),
                    SparseFeat('pos_sales_channel_id',2,embedding_dim=4,use_hash=False),
                    DenseFeat('pos_age', 1),
                    #DenseFeat('word_feature', 13312)
                    ]
    feature_columns += [
        VarLenSparseFeat(SparseFeat('hist_pos_artid',105542+1,embedding_dim=64,use_hash=False,embedding_name='pos_artid'),maxlen = 200,length_name="seq_length"),
        VarLenSparseFeat(SparseFeat('hist_pos_pcode',105542+1,embedding_dim=64,use_hash=False,embedding_name='pos_pcode'),maxlen = 200,length_name="seq_length")
    ]
    behavior_feature_list=["pos_artid","pos_pcode"]
    # if(use_neg):
    #     feature_columns+=[
    #         VarLenSparseFeat(SparseFeat('neg_hist_pos_artid',105542+1,embedding_dim=64,use_hash=False,embedding_name='pos_artid'),maxlen = 200,length_name="seq_length"),
    #         VarLenSparseFeat(SparseFeat('neg_hist_pos_pcode',105542+1,embedding_dim=64,use_hash=False,embedding_name='pos_pcode'),maxlen = 200,length_name="seq_length") 
    #     ]

    return feature_columns,behavior_feature_list

def get_xy(batch, feature_columns):
    #print(get_feature_names(feature_columns))
    permutation = np.random.permutation(batch['pos_custom_id'].shape[0]) 
    x = {name:batch[name][permutation] for name in get_feature_names(feature_columns)}
    y = batch['pos_click'][permutation]
    #print(y.shape)
    return x, y
    
def set_model(device = torch.device('cuda')):
    feature_columns, behavior_feature_list = get_fd()
    model=DIEN(feature_columns,behavior_feature_list,gru_type="AIGRU",device=device,use_negsampling=False)
    model.compile('adam','binary_crossentropy',metrics=['binary_crossentropy','auc'])
    return model, feature_columns, behavior_feature_list

def train(batchnum, model, feature_columns):
    
    #model.load_state_dict(torch.load("./Model/model.pt"))
    sampler = batch_getter(batchnum, use_neg=True, neg_sample=1)
    cnt = 0
    while(True):
        
        batch = sampler.get_batch(500)
        if batch['length'] == 0 :
            break
        print('batch getted')
        x, y = get_xy(batch, feature_columns)
        #print(x,y)
        model.fit(x, y, batch_size= 200, epochs=1, verbose=1, validation_split=0.2, shuffle=False)
        cnt+=500
        if(cnt % 10000 == 0):
            print(f'trained on {cnt} samples')
            torch.save(model.state_dict(), './Model/model4.pt')
        del batch
        del x
        del y
    torch.save(model.state_dict(), './Model/model4.pt')
    del sampler

# def validate(model, feature_columns):
#     sampler = batch_getter(31, use_neg=True, neg_sample=50)
#     cnt = 0
#     while(True):
        
#         batch = sampler.get_batch(100)
#         if batch['length'] == 0 :
#             break
#         print('batch getted')
#         x, y = get_xy(batch, feature_columns)
#         model.fit(x, y, batch_size=1500, epochs=10, verbose=1, validation_split=1, shuffle=True)
#         cnt+=100
#         if(cnt % 100000 == 0):
#             print(f'trained on {cnt} samples')
#         batch = None
#         x = None
#         y = None
#     #torch.save(model.state_dict(), './Model/model.pt')
#     del sampler


if __name__ == '__main__':
    device = torch.device('cuda')
    model, feature_columns, behavior_feature_list = set_model()
    for i in range(31):
        train(i, model, feature_columns)
        print(f'***********************************************file {i} trained ok!*****************************************************')
        








